package com.zzy.mccjt.retelimiter;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.zzy.mccjt.LimiterUtil;
import com.zzy.mccjt.retelimiter.exception.RateLimiterException;
import com.zzy.mccjt.retelimiter.method.ExpressionUtil;
import com.zzy.mccjt.retelimiter.method.LimiterConstants;
import com.zzy.mccjt.retelimiter.method.LimiterContext;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.EnableAspectJAutoProxy;
import org.springframework.stereotype.Component;
import org.springframework.util.DigestUtils;

import java.lang.reflect.Method;
import java.util.HashSet;
import java.util.Set;

/**
 * @program: mccjt-cloud
 * @description: 限流拦截器
 * @author: zhouyang.zheng
 * @create: 2022-08-01 16:00
 **/
@Slf4j
@Aspect
//@Component
public class RateLimiterInterceptor {
    private final RateLimiterTemplate rateLimiterTemplate;

    private final RateLimiterProperties rateLimiterProperties;

    private ObjectMapper objectMapper = new ObjectMapper();

    public RateLimiterInterceptor(RateLimiterTemplate rateLimiterTemplate, RateLimiterProperties rateLimiterProperties) {
        this.rateLimiterTemplate = rateLimiterTemplate;
        this.rateLimiterProperties = rateLimiterProperties;
    }

    @Pointcut(value = "@annotation(com.zzy.mccjt.retelimiter.RateLimiter)")
    public void cutService() {
        String a = "";
    }

    /**
     * 对限流注解处理
     * @param point
     * @return
     * @throws Throwable
     */
    @Around("cutService()")
    public Object interceptor(ProceedingJoinPoint point) throws Throwable {
        if (!rateLimiterProperties.getEnabled()) {
            return point.proceed();
        }
        boolean allow = true;
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        RateLimiter rateLimiter = method.getAnnotation(RateLimiter.class);
        if (rateLimiter != null) {
            allow = checkLimit(rateLimiter, point, new HashSet<>(1));
        } else {
            RateLimiters rateLimiters = method.getAnnotation(RateLimiters.class);
            if (rateLimiters != null) {
                Set<String> keyCache = new HashSet<>(rateLimiters.value().length);
                for (RateLimiter limiter : rateLimiters.value()) {
                    boolean checkResult = checkLimit(limiter, point, keyCache);
                    if (!checkResult) {
                        allow = false;
                    }
                }
            }
        }

        if (!allow) {
            throw new RateLimiterException("访问频率超过限制");
        } else {
            return point.proceed();
        }
    }

    private boolean checkLimit(RateLimiter rateLimiter, ProceedingJoinPoint point, Set<String> keyCache) {
        boolean allow = true;
        LimiterType limitType = null;
        String key = "";
        try {
            MethodSignature signature = (MethodSignature) point.getSignature();
            Method method = signature.getMethod();
            limitType = rateLimiter.type();
            switch (limitType) {
                case PATH:
                    key = LimiterUtil.getPathKey(signature);
                    if (log.isDebugEnabled()) {
                        log.debug("limiter by path, path hash code: {}", key);
                    }
                    break;
                case IP:
                    key = LimiterUtil.getIpAddress();
                    if (log.isDebugEnabled()) {
                        log.debug("limiter by ip, ip: {}", key);
                    }
                    break;
                case KEY:
                    StringBuilder keySb = new StringBuilder(LimiterUtil.getPathKey(signature));
                    keySb.append(":");
                    if (LimiterConstants.isUndefined(rateLimiter.key())) {
                        byte[] argsJson = objectMapper.writeValueAsBytes(point.getArgs());
                        keySb.append(DigestUtils.md5DigestAsHex(argsJson));
                    } else {
                        LimiterContext limiterContext = LimiterContext.builder()
                                .args(point.getArgs())
                                .method(method)
                                .build();
                        keySb.append(ExpressionUtil.evalKey(limiterContext, rateLimiter.key(), method));
                    }
                    key = keySb.toString();
                    if (log.isDebugEnabled()) {
                        log.debug("limiter by key, scriptKey: {}, key: {}", rateLimiter.key(), key);
                    }
                    break;
                default:
                    break;
            }
            if (!keyCache.contains(key)) {
                RateLimitCondition condition = RateLimitCondition.builder()
                        .key(key)
                        .burstCapacity(rateLimiter.burstCapacity())
                        .replenishRate(rateLimiter.replenishRate())
                        .requestedTokens(rateLimiter.requestedTokens())
                        .build();
                allow = rateLimiterTemplate.check(condition);
                keyCache.add(key);
            } else {
                log.warn("访问限流存在重复配置, limit type: {}, key: {}", limitType, key);
            }
        } catch (Exception e) {
            log.error(String.format("访问限流异常,请检查配置, limit type: %s, key: %s", limitType, key), e);
        }
        return allow;
    }
}
